In [16]:

    

from os import path, mkdir
from subprocess import call
from time import time
import numpy as np
import pandas as pd
import yaml
from nbodykit.cosmology import Cosmology
from nbodykit.cosmology.power import EHPower, ZeldovichPower


    

In [17]:

    

from pearce.sims import sim_config_template, submission_file_template


    

In [25]:

    

def submit_sims(lhc_fname, output_dir, powerspectrum, initial_z, final_z,\
                npart, boxsize, ncores, sim_time):

    assert path.isdir(output_dir), "Invalid output directory %s"%output_dir

    lhc = pd.read_csv('/home/users/swmclau2/Git/pearce/bin/pmesh_tests/'+lhc_fname, sep=' ', index_col=None)
    param_names = list(lhc)

    k = np.logspace(-4, 2, 1000)
    # others?
    assert powerspectrum.lower() in {'linear', 'zeldovich'}, "Invalid initial power specturm"

    if powerspectrum.lower() == 'linear':
        power = EHPower
    else:
        power = ZeldovichPower

    initial_a, final_a = 1./(1+initial_z), 1./(1+final_z)

    for idx, row in lhc.iterrows():
        # TODO what to do if already exists? likely will throw some kind of error
        
        boxdir = path.join(output_dir, "Box_%03d/"%idx)
        #mkdir(boxdir)

        # TODO parse Jeremy's header, or adjust it to be compliant here
        param_dict = dict([(pn, row[pn]) for pn in param_names] )
        print param_dict.keys()
        param_dict['N_ncdm'] = 3.0
        param_dict['N_ur'] = param_dict['Neff']-3.0
        del param_dict['Neff']
        
        param_dict['h'] = param_dict['H0']/100
        del param_dict['H0']
        
        param_dict['w0_fld'] = param_dict['w0']
        del param_dict['w0']
        
        cosmo = Cosmology(**param_dict)
        # TODO i wanted to do this with fixed amplitudes, not sure how
        p = power(cosmo, initial_z)(k)
        np.savetxt(path.join(boxdir,'powerspec.txt'), np.array((k,p)).T)
        return 0
        with open(path.join(boxdir, 'nbodykit.lua')) as f:
           f.write(sim_config_template.format(nc = npart, boxsize = boxsize,
                                              initial_a = initial_a,
                                              final_a = final_a,
                                              final_z = final_z,
                                              omega_m = cosmo.Om0,
                                              h = cosmo.h,
                                              seed = int(time())%1000,
                                              boxdir = boxdir))

        # now, write submission file and submit

        with open(path.join(boxdir, 'submit_sim.sbatch')) as f:
            f.write(submission_file_template.format(jobname = 'Box_%03d'%idx,
                                                    boxdir = boxdir,
                                                    time = sim_time*60, #TODO in config
                                                    ntasks = ncores))

        call("sbatch {boxdir}submit_sim.sbatch".format(boxdir), shell=True)


    

In [26]:

    

config_fname = '/home/users/swmclau2/Git/pearce/bin/pmesh_tests/pmesh_test.yaml'
# NOTE could use a vars(args) to turn into a dictionary.
# currently not doing that if some of those end up not being arguments for submit_sims
try:
    assert path.isfile(config_fname)
except AssertionError:
    raise AssertionError("%s is not a valid filename." % config_fname)

with open(config_fname, 'r') as ymlfile:
    cfg = yaml.load(ymlfile)

#try:
submit_sims(**cfg)
#except TypeError:
#    raise TypeError("Not all require arguments specified in yaml config.")


    

    




['Neff', 'H0', 'w0', 'omch2', 'ln10As', 'ns', 'ombh2']






    
Out[26]:





0

In [ ]: